

package com.hazelcast.jet.impl.pipeline.transform;

import com.hazelcast.function.FunctionEx;
import com.hazelcast.function.PredicateEx;
import com.hazelcast.jet.impl.pipeline.PipelineImpl.Context;
import com.hazelcast.jet.impl.pipeline.Planner;
import com.hazelcast.jet.impl.pipeline.Planner.PlannerVertex;

import javax.annotation.Nonnull;
import java.io.Serial;

import static com.hazelcast.jet.core.Vertex.LOCAL_PARALLELISM_USE_DEFAULT;
import static com.hazelcast.jet.core.processor.DiagnosticProcessors.peekOutputP;

public class PeekTransform<T> extends AbstractTransform {

    @Serial
    private static final long serialVersionUID = 1L;

    @Nonnull
    public final PredicateEx<? super T> shouldLogFn;
    @Nonnull
    public final FunctionEx<? super T, ? extends CharSequence> toStringFn;

    public PeekTransform(@Nonnull Transform upstream, @Nonnull PredicateEx<? super T> shouldLogFn, @Nonnull FunctionEx<? super T, ? extends CharSequence> toStringFn) {
        super("peek", upstream);
        this.shouldLogFn = shouldLogFn;
        this.toStringFn = toStringFn;
    }

    @Override
    public void addToDag(Planner p, Context context) {
        determineLocalParallelism(LOCAL_PARALLELISM_USE_DEFAULT, context, p.isPreserveOrder());
        PlannerVertex peekedPv = p.transform2vertex.get(this.upstream().get(0));
        // Peeking transform doesn't add a vertex, so point to the upstream
        // transform's vertex:
        p.transform2vertex.put(this, peekedPv);
        peekedPv.v.updateMetaSupplier(sup -> peekOutputP(toStringFn, shouldLogFn, sup));
    }
}
